#include "gpu.h"
#include <cuda_runtime.h>
#include <iostream>
#include <cstring>

// 全局变量声明（已在guessing.cpp中定义）
int total_guesses = 0;

// CUDA内核函数：处理单段PT的猜测生成
__global__ void processStringKernel(const char* prefix, int prefix_len, 
                                    const char** values, int* lengths, int num,
                                    char* output, int* output_offsets) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < num) {
        int val_len = lengths[idx];
        int result_len = prefix_len + val_len;
        int offset = output_offsets[idx];
        
        // 复制前缀字符串
        for (int i = 0; i < prefix_len; i++) {
            output[offset + i] = prefix[i];
        }
        
        // 复制值字符串
        for (int i = 0; i < val_len; i++) {
            output[offset + prefix_len + i] = values[idx][i];
        }
        
        // 添加终止符
        output[offset + result_len] = '\0';
    }
}

void GPU(std::vector<std::string> &val, int num, std::vector<std::string> &guesses, const std::string &guess) {
    if (num <= 0 || val.empty()) return;

    // 准备主机端数据
    const char* prefix = guess.c_str();
    int prefix_len = guess.length();
    
    // 创建值的C风格字符串数组
    const char** h_values = new const char*[num];
    int* h_lengths = new int[num];
    int* h_offsets = new int[num];
    
    int total_len = 0;
    for (int i = 0; i < num; i++) {
        h_values[i] = val[i].c_str();
        h_lengths[i] = val[i].length();
        h_offsets[i] = total_len;
        total_len += prefix_len + h_lengths[i] + 1; // +1 for null terminator
    }
    
    // 分配设备内存
    char* d_prefix;
    const char** d_values;
    int* d_lengths;
    int* d_offsets;
    char* d_output;
    
    cudaMalloc(&d_prefix, prefix_len + 1);
    cudaMalloc(&d_values, num * sizeof(char*));
    cudaMalloc(&d_lengths, num * sizeof(int));
    cudaMalloc(&d_offsets, num * sizeof(int));
    cudaMalloc(&d_output, total_len * sizeof(char));
    
    // 为每个值分配内存并复制数据
    const char** temp_values = new const char*[num];
    for (int i = 0; i < num; i++) {
        char* d_value;
        int len = h_lengths[i] + 1; // +1 for null terminator
        cudaMalloc(&d_value, len);
        cudaMemcpy(d_value, h_values[i], len, cudaMemcpyHostToDevice);
        temp_values[i] = d_value;
    }
    
    // 复制数据到设备
    cudaMemcpy(d_prefix, prefix, prefix_len + 1, cudaMemcpyHostToDevice);
    cudaMemcpy(d_values, temp_values, num * sizeof(char*), cudaMemcpyHostToDevice);
    cudaMemcpy(d_lengths, h_lengths, num * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_offsets, h_offsets, num * sizeof(int), cudaMemcpyHostToDevice);
    
    // 启动内核
    int blockSize = 256;
    int gridSize = (num + blockSize - 1) / blockSize;
    processStringKernel<<<gridSize, blockSize>>>(d_prefix, prefix_len, d_values, d_lengths, num, d_output, d_offsets);
    
    // 同步设备
    cudaDeviceSynchronize();
    
    // 将结果复制回主机
    char* h_output = new char[total_len];
    cudaMemcpy(h_output, d_output, total_len, cudaMemcpyDeviceToHost);
    
    // 构建结果字符串
    for (int i = 0; i < num; i++) {
        guesses.push_back(std::string(h_output + h_offsets[i]));
    }
    
    // 更新全局计数
    total_guesses += num;
    
    // 释放内存
    for (int i = 0; i < num; i++) {
        cudaFree((void*)temp_values[i]);
    }
    delete[] temp_values;
    delete[] h_values;
    delete[] h_lengths;
    delete[] h_offsets;
    delete[] h_output;
    
    cudaFree(d_prefix);
    cudaFree(d_values);
    cudaFree(d_lengths);
    cudaFree(d_offsets);
    cudaFree(d_output);
}